
import os
import cv2
from qtpy import QtWidgets
from qtpy import QtCore
from qtpy import QtGui

from qtpy.QtCore import Qt

# 自定义模块

from qt import add_actions
from qt import create_action
from qt import create_menu
from qt import getCurrentFolderImages

from format import FreeLabelFormat
from format import YoloLableFormat
from utils import create_struct

from widgets import DrawType
from widgets import ToolBar
from widgets import Canvas
from widgets import ImageSlider
from widgets import LabelDialog
from widgets import LabelManageDialog
from widgets import LabelAttrDialog
from widgets import LabelFormat

from label import Shape

from enum import Enum


class MainWindow(QtWidgets.QMainWindow):

    def __init__(self) -> None:
        super(MainWindow, self).__init__()
        # 定义文件列表
        self.file_list = []
        self.label_info = None
        self.current_image = None
        # 标签保存
        self.labelSavePath = None
        self.LabelFormat = FreeLabelFormat()

        self.labelFormaAction = QtWidgets.QWidgetAction(self)
        self.labelFormatComboBox = LabelFormat(self)
        self.labelFormaAction.setDefaultWidget(self.labelFormatComboBox)

        self.imageSliderAction = QtWidgets.QWidgetAction(self)
        self.imageSlider = ImageSlider(self)
        self.imageSliderAction.setDefaultWidget(self.imageSlider)

        # 是否需要保存
        self.isSave = False

        # 恢复设置
        self.settings = QtCore.QSettings("标注", "自由标注")
        self.restoreSettings()
        # 获取主菜单
        self.allActions = create_struct()

        self.menu = self.createMainMenu()
        # 创建文件菜单
        self.createFileMenu()
        self.createEditMenu()
        

        # 编辑标签对话框
        self.label_dlg = LabelDialog(self)

        self.labelattr_dlg = LabelAttrDialog(self)
        self.labelattr_dockwidet = QtWidgets.QDockWidget(self.tr("标注信息"), self)
        self.labelattr_dockwidet.setObjectName(self.tr("标注信息"))  # QMainWindow::saveState(): 'objectName' not set for QDockWidget 0x27024d0 '标注信息;
        self.labelattr_dockwidet.setAllowedAreas(Qt.RightDockWidgetArea)
        self.labelattr_dockwidet.setWidget(self.labelattr_dlg)
        self.labelattr_dockwidet.setFloating(False)
        self.addDockWidget(Qt.RightDockWidgetArea, self.labelattr_dockwidet)
        # 工具栏
        # 标注工具
        self.drawToolStatus = DrawType.DRAW_TYPE_NONE.value
        # 创建标注工具栏
        self.create_draw_toolbar()
        # 创建标签编辑工具栏
        self.create_label_toolbar()
        # 创建图片变换的工具栏
        self.create_image_tf_toolbar()


        # 创建图片操作工具栏
        self.create_image_view_toolbar()
        # 查看图片工具
        # self.imageViewTool = self.createToolbar("ImageView")
        # self.create_image_view_toolbar()

        self.canvas = Canvas()

        scrollArea = QtWidgets.QScrollArea()
        scrollArea.setWidget(self.canvas)
        scrollArea.setWidgetResizable(True)
        self.setCentralWidget(scrollArea)
        # 最后打开的文件夹
        self.lastOpenDir = None

        # 连接信号槽
        self.connectSignal()

        # 定义状态栏
        self.status_messages = QtWidgets.QLabel("")
        self.statusBar().addPermanentWidget(self.status_messages)



    def connectSignal(self):
        # 图片缩放
        self.canvas.zoomChangeSignal.connect(self.zoomChange)
        self.canvas.drawCompletedSignal.connect(self.drawCompleted)
        self.canvas.modifyShapeSignal.connect(self.setIsSave)
        self.imageSlider.valueChanged.connect(self.showImage)
        self.labelFormatComboBox.currentIndexChanged.connect(self.labelFormatChange)

    def labelFormatChange(self, index):
        if index == 0:
            self.LabelFormat = FreeLabelFormat()
        elif index == 1:
            self.LabelFormat = YoloLableFormat()
        else:
            pass
        # TODO 

    def restoreSettings(self):
        """恢复设置"""
        # 获取设置信息
        size = self.settings.value("window/size", QtCore.QSize(600, 500))
        position = self.settings.value("window/position", QtCore.QPoint(0, 0))
        state = self.settings.value("window/state", QtCore.QByteArray())
        # 恢复设置
        self.resize(size)
        self.move(position)
        self.restoreState(state)

    def storeSettings(self):
        """保存设置信息"""
        self.settings.setValue("window/size", self.size())
        self.settings.setValue("window/position", self.pos())
        self.settings.setValue("window/state", self.saveState())

    def createFileMenu(self):
        """创建文件菜单"""
        file_menu = create_struct()
        # 打开文件
        file_menu.open = create_action(
            parent_widget=self,
            text=self.tr("打开"),
            slot=self.openFile,
            tip=self.tr("打开图片")
        )
        # 打开文件夹
        file_menu.opendir = create_action(
            parent_widget=self,
            text=self.tr("打开文件夹"),
            slot=self.openDir,
            icon="folder-open",
            tip=self.tr("打开文件夹")
        )
        # 保存当前图片的标签
        file_menu.save_label = create_action(
            parent_widget=self,
            text=self.tr("保存"),
            slot=self.saveLabelFile,
            icon="save",
            tip=self.tr("保存")
        )
        # 自动保存当前图片的标签
        file_menu.auto_save_label = create_action(
            parent_widget=self,
            text=self.tr("自动保存"),
            slot=None,
            tip=self.tr("自动保存"),
            checkable=True,
            checked=False
        )

        # 退出程序
        file_menu.quit = create_action(
            parent_widget=self,
            text=self.tr("退出"),
            slot=self.close,
            tip=self.tr("退出程序")
        )
        # 添加到文件菜单中
        add_actions(
            self.menu.file,
            (
            file_menu.open,
            file_menu.opendir,
            None,
            file_menu.save_label,
            file_menu.auto_save_label,
            None,
            file_menu.quit,
            )
        )
        # 保存 action对象
        self.allActions.file_menu = file_menu

    def createEditMenu(self):
        """创建编辑菜单"""

        edit_menu = create_struct()

        edit_menu.label_manager = create_action(
            parent_widget=self,
            text=self.tr("标签管理"),
            slot=self.labelManager,
            tip=self.tr("打开标签管理界面")
        )
        # 添加到文件菜单中
        add_actions(
            self.menu.edit,
            (
            edit_menu.label_manager,
            )
        )
        # 保存 action对象
        self.allActions.edit_menu = edit_menu

    def createViewMenu(self):
        """创建视图菜单"""
        # TODO something

    def create_help_menu(self):
        """创建帮助菜单"""

    def create_image_view_toolbar(self):
        """创建查看工具"""
        self.image_view_toolbar = self.createToolbar(self.tr("查看"), Qt.ToolButtonTextUnderIcon, Qt.TopToolBarArea)
        image_view_actions = create_struct()

        image_view_actions.first_image = create_action(
            parent_widget=self,
            text=self.tr("开头"),
            slot=self.firstImage,
            icon="first",
            enabled=False
        )
        image_view_actions.pre_image = create_action(
            parent_widget=self,
            text=self.tr("向前"),
            slot=self.preImage,
            shortcut="A",
            icon="pre",
            enabled=False
        )
        image_view_actions.pause_image = create_action(
            parent_widget=self,
            text=self.tr("播放"),
            slot=self.pauseImage,
            icon="play",
            enabled=False
        )
        image_view_actions.next_image = create_action(
            parent_widget=self,
            text=self.tr("向后"),
            slot=self.nextImage,
            shortcut="D",
            icon="next",
            enabled=False
        )
        image_view_actions.last_image = create_action(
            parent_widget=self,
            text=self.tr("末尾"),
            slot=self.lastImage,
            icon="last",
            enabled=False
        )
        add_actions(
            self.image_view_toolbar,
            (
            image_view_actions.first_image,
            image_view_actions.pre_image,
            image_view_actions.pause_image,
            image_view_actions.next_image,
            image_view_actions.last_image,
            self.imageSliderAction
            )
        )
        self.allActions.image_view_actions = image_view_actions

    def create_draw_toolbar(self):
        """创建标注工具"""
        self.draw_toolbar = self.createToolbar("标注", Qt.ToolButtonTextUnderIcon, Qt.LeftToolBarArea)
        draw_tools_actions = create_struct()

        # draw_tools_actions.draw_point = create_action(
        #     parent_widget=self,
        #     text=self.tr("点"),
        #     slot=lambda: self.updateDrawToolActionStatus(DrawType.DRAW_TYPE_POINT.value),
        #     icon="dot"
        # )
        draw_tools_actions.draw_rect = create_action(
            parent_widget=self,
            text=self.tr("框"),
            slot=lambda: self.updateDrawToolActionStatus(DrawType.DRAW_TYPE_RECT.value),
            icon="rect"
        )
        # TODO 线

        # draw_tools_actions.draw_polygon = create_action(
        #     parent_widget=self,
        #     text=self.tr("多边形"),
        #     slot=lambda: self.updateDrawToolActionStatus(DrawType.DRAW_TYPE_POLYGON.value),
        #     icon="polygon"
        # )
        add_actions(
            self.draw_toolbar,  
            (
            # draw_tools_actions.draw_point,
            draw_tools_actions.draw_rect,
            # draw_tools_actions.draw_polygon,
            )
        )
        self.allActions.draw_tools_actions = draw_tools_actions

    def create_image_tf_toolbar(self):
        self.img_tf_toolbar = self.createToolbar(self.tr("变换"), Qt.ToolButtonTextUnderIcon, Qt.TopToolBarArea)
        img_tf_actions = create_struct()

        img_tf_actions.pan = create_action(
            parent_widget=self,
            text=self.tr("平移"),
            slot=None,
            icon="pan",
            enabled=False
        )
        img_tf_actions.undo = create_action(
            parent_widget=self,
            text=self.tr("恢复"),
            slot=self.unDo,
            icon="undo",
            enabled=False
        )
        img_tf_actions.zoom_in = create_action(
            parent_widget=self,
            text=self.tr("放大"),
            slot=self.zoomIn,
            icon="zoom_in",
            enabled=False
        )
        img_tf_actions.zoom_out = create_action(
            parent_widget=self,
            text=self.tr("缩小"),
            slot=self.zoomOut,
            icon="zoom_out",
            enabled=False
        )
        add_actions(
            self.img_tf_toolbar,
            (
                img_tf_actions.pan,
                img_tf_actions.undo,
                img_tf_actions.zoom_in,
                img_tf_actions.zoom_out,
            )
        )

        self.allActions.img_tf_actions = img_tf_actions

    def create_label_toolbar(self):
        self.label_toolbar = self.createToolbar(self.tr("编辑标注"), Qt.ToolButtonTextUnderIcon, Qt.TopToolBarArea)
        label_actions = create_struct()

        label_actions.pick = create_action(
            parent_widget=self,
            text=self.tr("拾取"),
            slot=self.pickObject,
            icon="pick"
        )

        label_actions.delete = create_action(
            parent_widget=self,
            text=self.tr("删除"),
            slot=self.deleteObject,
            shortcut="R",
            icon="delete",
            enabled=True
        )

        add_actions(
            self.label_toolbar,
            (
                label_actions.pick,
                label_actions.delete,
                self.labelFormaAction,
            )
        )
        self.allActions.label_actions = label_actions

    def createMainMenu(self):
        """创建主菜单"""
        menu = create_struct(
            file=create_menu(self, self.tr("&文件")),
            edit=create_menu(self, self.tr("&编辑")),
            view=create_menu(self, self.tr("&视图")),
            help=create_menu(self, self.tr("&帮助")),
        )
        return menu

    def createToolbar(self, title, toolButtonStyle, area, actions=None):
        # 创建工具栏函数
        toolbar = ToolBar(title)
        toolbar.setObjectName("%sToolBar" % title)
        toolbar.setToolButtonStyle(toolButtonStyle)
        if actions:
            add_actions(toolbar, actions)
        self.addToolBar(area, toolbar)
        return toolbar

    # 菜单选项实现函数
    def openFile(self):

        if not self.mayContinue():
            return

        from PyQt5.QtWidgets import QFileDialog
        dir = QFileDialog()     # 创建文件对话框
        dir.setFileMode(QFileDialog.ExistingFiles)     # 设置多选
        dir.setDirectory(".")     # 设置初始路径为D盘
        # 设置只显示图片文件
        dir.setNameFilter("图片文件(*.jpg *.png *.bmp *.ico *.gif)")
        file_list = []
        if dir.exec_():      # 判断是否选择了文件
            file_list = dir.selectedFiles()
        if len(file_list) == 0:
            return
        self.imageSlider.setHSliderValue(0)

    def openDir(self):
        """打开文件夹"""

        if not self.mayContinue():
            return

        if self.lastOpenDir and os.path.exists(self.lastOpenDir):
            open_path = self.lastOpenDir
        else:
            open_path = "."
        targetDirPath = str(
            QtWidgets.QFileDialog.getExistingDirectory(
                self,
                self.tr("打开文件夹"),
                open_path,
                QtWidgets.QFileDialog.ShowDirsOnly | QtWidgets.QFileDialog.DontResolveSymlinks,
            )
        )
        self.lastOpenDir = targetDirPath
        self.loadCategoriesFile()
        self.file_list = getCurrentFolderImages(targetDirPath)
        if self.file_list:
            self.imageSlider.setMaximum(len(self.file_list) - 1)
            self.imageSlider.setHSliderValue(0)
        else:
            pass

    def writeCategoriesFile(self):
        """写入categories"""
        import json
        if self.lastOpenDir:
            self.LabelFormat.writeCategoriesFile(self.lastOpenDir, self.label_info)

    def updateEditLabelDlgLabelInfo(self, clear=False):
        if clear:
            self.label_dlg.clearInfo()
        self.label_dlg.setLabelInfo(self.label_info)

    def loadCategoriesFile(self):
        if self.lastOpenDir:
            label_info = self.LabelFormat.readCategoriesFile(self.lastOpenDir)
            if label_info is not None:
                self.label_info = label_info
                self.updateEditLabelDlgLabelInfo(clear=True)

    def labelManager(self):
        """标签管理"""
        if self.lastOpenDir:
            label_manager = LabelManageDialog(self)
            label_manager.loadLabel(self.label_info)
            new_label_info = None
            if label_manager.exec_():
                new_label_info = label_manager.getLabelData()
            else:
                return
            self.label_info = new_label_info
            self.updateEditLabelDlgLabelInfo()
            self.writeCategoriesFile()

    def pickObject(self):
        self.canvas.setEditing(True)

    def setIsSave(self):
        if not self.isSave:
            self.isSave = True

    def deleteObject(self):
        ret = self.canvas.deleteSelected()
        self.updateLabelTreeList()
        if ret:
            self.setIsSave()

    def zoomIn(self):
        """放大图片"""
        if self.current_image is not None:
            self.canvas.zoom_in()

    def zoomOut(self):
        """缩小图片"""
        if self.current_image is not None:
            self.canvas.zoom_out()

    def unDo(self):
        """恢复图片默认比例"""
        if self.current_image is not None:
            self.canvas.setScale(1.0)

    def firstImage(self):
        """跳转到第一张图片"""
        if not self.mayContinue():
            return
        if self.file_list and self.current_image_id != 0:
            self.imageSlider.setHSliderValue(0)
    def preImage(self):
        """跳转到前一张图片"""
        if not self.mayContinue():
            return
        if self.file_list and self.current_image_id != 0:
            self.imageSlider.setHSliderValue(self.current_image_id - 1)

    def pauseImage(self):
        """播放暂停图片"""

    def nextImage(self):
        """跳转到下一张图片"""
        if not self.mayContinue():
            return
        if self.file_list and self.current_image_id != len(self.file_list) - 1:
            self.imageSlider.setHSliderValue(self.current_image_id + 1)

    def lastImage(self):
        """跳转到最后一张图片"""
        if not self.mayContinue():
            return
        if self.file_list and self.current_image_id != len(self.file_list) - 1:
            self.imageSlider.setHSliderValue(len(self.file_list) - 1)


    # 工具栏 动作函数
    def getDrawToolActionStatus(self, status):
        return (self.drawToolStatus & status)

    def get_current_labelinfo(self):
        # 弹出选择标签对话框
        self.canvas.current_labelinfo = self.labelDlgPopUp()
        if self.canvas.current_labelinfo is not None:
            self.canvas.current_labelinfo['color'] = self.LabelFormat.get_label_color(self.canvas.current_labelinfo['name'])
            return True
        else:
            return False

    def updateDrawToolActionStatus(self, status):

        if self.get_current_labelinfo():
            self.drawToolStatus = status
            self.canvas.setEditing(False)
            self.canvas.setDrawingType(status)

        # self.allActions.draw_tools_actions.draw_point.setEnabled(not self.getDrawToolActionStatus(DrawType.DRAW_TYPE_POINT.value))
        # self.allActions.draw_tools_actions.draw_rect.setEnabled(not self.getDrawToolActionStatus(DrawType.DRAW_TYPE_RECT.value))
        # self.allActions.draw_tools_actions.draw_polygon.setEnabled(not self.getDrawToolActionStatus(DrawType.DRAW_TYPE_POLYGON.value))

    # 系统事件
    def closeEvent(self, event):
        # 关闭时保存设置信息
        self.storeSettings()

    # 槽函数
    def zoomChange(self, delta, pos):
        """缩放图片"""
        pass

    def labelDlgPopUp(self, move=True):
        if move:
            # 移动窗口到鼠标位置
            self.label_dlg.move(QtGui.QCursor.pos())
        if self.label_dlg.exec_():
            # 确认按钮
            return self.label_dlg.getCurrentLabelInfo()
        return None

    def updateLabelTreeList(self):
        self.labelattr_dlg.clearShapes()
        self.labelattr_dlg.addLabelTree(self.canvas.shape_list)

    def drawCompleted(self):
        # 完成标注
        # 撤销存在问题
        # if self.current_labelinfo is not None:
            self.updateLabelTreeList()
            if self.allActions.file_menu.auto_save_label.isChecked():
                # 自动保存
                label_file = os.path.splitext(self.get_current_filename())[0] + self.LabelFormat.suffix
                self.saveLabelData(label_file)
            else:
                # 设置需要保存的状态为真
                self.isSave = True
        # else:
        #     # 撤销操作，恢复最后一步操作
        #     self.canvas.undoLastLine()

    def saveFileDialog(self):
        """
        保存标签对话框，获取保存路径
        """
        # 定义筛选条件
        filters = self.tr("Label files (*{})".format(self.LabelFormat.suffix))
        current_filename = self.get_current_filename()
        # 定义打开文件对话框
        dlg = QtWidgets.QFileDialog(
            self,
            self.tr("保存路径"),
            os.path.dirname(current_filename),
            filters
        )
        # 获取文件名
        filename = os.path.splitext(os.path.basename(current_filename))[0]
        # 获取文件路径
        filedir = os.path.dirname(current_filename)
        # 定义标签文件路径
        default_save_filename = os.path.join(filedir, filename + self.LabelFormat.suffix)
        filename = dlg.getSaveFileName(
            self,
            self.tr("保存路径"),
            default_save_filename,
            filters
        )

        if isinstance(filename, tuple):
            filename, _ = filename
        return filename

    def saveLabelFile(self):
        """
        保存标签文件，如果文件已经存在则不需要弹出保存对话框
        """
        self.saveLabelData(self.saveFileDialog())

    def saveLabelData(self, save_path):
        """
        保存标签数据
        """
        def format_shape(shape):
            data = dict()
            data.update(
                dict(
                    label=shape.label,
                    points=[(p.x(), p.y()) for p in shape.point_list],
                    group_id=shape.group_id,
                    attributes=shape.attributes,
                    shape_type=shape.shape_type
                )
            )
            return data
        if save_path == "":
            return
        shapes = [format_shape(item) for item in self.canvas.shape_list]
        image_w = self.current_image.width()
        image_h = self.current_image.height()
        self.LabelFormat.write(save_path, shapes, image_w, image_h)

        self.resetStates()

    def loadShapps(self, shapes):
        ret = []
        for shape in shapes:
            label = shape["label"]
            shape_type = shape["shape_type"]
            # group_id = shape["group_id"]
            attributes = shape["attributes"]
            point_list = shape["points"]

            if "color" not in shape or shape["color"] is None:
                shape = Shape(label=label, shape_type=shape_type, attributes=attributes)
            else:
                color = eval(shape["color"])
                shape = Shape(label=label, color=color, shape_type=shape_type, attributes=attributes)
            for x, y in point_list:
                shape.addPoint(QtCore.QPointF(x, y))
            shape.setClose()
            ret.append(shape)
        self.canvas.loadShapes(ret)
        self.updateLabelTreeList()

    def loadFile(self, filename):

        # 创建标签文件名称
        label_file = os.path.splitext(filename)[0] + self.LabelFormat.suffix

        # 设置状态栏
        self.status(self.tr("加载{}...".format(os.path.basename(filename))))
        # 加载图片
        pixmap = QtGui.QPixmap(filename)
        self.current_image = pixmap.toImage()
        self.canvas.setPixmap(pixmap)
        # 清除显示列表
        self.updateLabelTreeList()
        # 加载标签
        if os.path.exists(label_file):
            self.LabelFormat.read(label_file, (pixmap.height(), pixmap.width()))
            self.loadShapps(self.LabelFormat.shapes)
        self.canvas.setScale(self.scaleFitWindow())

        # 设置状态栏
        self.status(self.tr("已经加载{}".format(os.path.basename(filename))))

    def showImage(self, current_image_id):
        """显示图片"""

        self.current_image_id = current_image_id
        file_path = self.get_current_filename()
        self.loadFile(file_path)

        if len(self.file_list) > 1:
            if self.current_image_id == 0:
                self.allActions.image_view_actions.first_image.setEnabled(False)
                self.allActions.image_view_actions.pre_image.setEnabled(False)
            else:
                self.allActions.image_view_actions.first_image.setEnabled(True)
                self.allActions.image_view_actions.pre_image.setEnabled(True)

            if self.current_image_id == len(self.file_list) - 1:
                self.allActions.image_view_actions.next_image.setEnabled(False)
                self.allActions.image_view_actions.last_image.setEnabled(False)
            else:
                self.allActions.image_view_actions.next_image.setEnabled(True)
                self.allActions.image_view_actions.last_image.setEnabled(True)

    def status(self, message, delay=5000):
        self.statusBar().showMessage(message, delay)

    def get_current_filename(self):
        return self.file_list[self.current_image_id]

    def mayContinue(self):
        if not self.isSave:
            return True
        if self.allActions.file_menu.auto_save_label.isChecked():
            # 自动保存
            label_file = os.path.splitext(self.get_current_filename())[0] + self.LabelFormat.suffix
            self.saveLabelData(label_file)
            return True
        mb = QtWidgets.QMessageBox
        msg = self.tr('Save annotations to "{}" before closing?').format(
            self.get_current_filename()
        )
        answer = mb.question(
            self,
            self.tr("Save annotations?"),
            msg,
            mb.Save | mb.Discard | mb.Cancel,
            mb.Save,
        )
        if answer == mb.Discard:
            self.resetStates()
            return True
        elif answer == mb.Save:
            label_file = os.path.splitext(self.get_current_filename())[0] + self.LabelFormat.suffix
            self.saveLabelData(label_file)
            return True
        else:  # answer == mb.Cancel
            return False

    def resetStates(self):
        # 重置状态
        self.isSave = False

    def scaleFitWindow(self):
        """Figure out the size of the pixmap to fit the main widget."""
        e = 2.0  # So that no scrollbars are generated.
        w1 = self.centralWidget().width() - e
        h1 = self.centralWidget().height() - e
        a1 = w1 / h1
        # Calculate a new scale value based on the pixmap's aspect ratio.
        w2 = self.canvas.pixmap.width() - 0.0
        h2 = self.canvas.pixmap.height() - 0.0
        a2 = w2 / h2
        return w1 / w2 if a2 >= a1 else h1 / h2
